
from ultralytics import YOLO
import streamlit as st
import cv2
from PIL import Image,ImageDraw,ImageFont
import tempfile
# import config
from src.infer import *
def _display_detected_frames(model, st_frame, image):
    img_rgb = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA))
    res_plotted = predict_image(model,img_rgb)
    st_frame.image(res_plotted,
                   caption='Detected Video',
                   channels="BGR",
                   use_column_width=True
                   )


@st.cache_resource
def load_model(model_path):
    model = YOLO(model_path)
    return model

def predict(model,image):
    res = model.predict(image)
    labels = res[0].names
    probs = res[0].probs.cpu().numpy()
    best_result = probs.top1
    confidence = probs.top1conf
    return best_result, confidence, labels[best_result]

def predict_image(model,uploaded_image):
    res = model.predict(uploaded_image)
    labels = res[0].names
    probs = res[0].probs.cpu().numpy()
    top1_index = probs.top1
    top1_confidence = probs.top1conf
    uploaded_image = uploaded_image.convert('RGB')
    draw = ImageDraw.Draw(uploaded_image)
    font = ImageFont.truetype("simhei.ttf", 64)
    class_name = labels[top1_index]
    if class_name=='normal':
        class_name='正常'
        if top1_confidence<0.7:
            class_name='状态异常'
    elif top1_confidence > 0.92:
        class_name='状态异常'
    else:
        class_name='正常'

    confidence = top1_confidence * 100
    text = '{:<5} {:>.2f}%'.format(class_name, confidence)
    draw.text((10, 10), text, font=font,fill=(255, 0, 0, 1))
    return uploaded_image
#图像推理
def infer_uploaded_image(model):
    source_img = st.file_uploader(
        label="选择一张图片...",
        type=("jpg", "jpeg", "png", 'bmp', 'webp')
    )

    col1, col2 = st.columns(2)

    with col1:
        if source_img:
            uploaded_image = Image.open(source_img)
            # adding the uploaded image to the page with caption
            st.image(
                image=source_img,
                caption="Uploaded Image",
                use_column_width=True
            )

    if source_img:
        if st.button("Execution"):
            with st.spinner("执行中..."):
                uploaded_image = predict_image(model, uploaded_image)
                with col2:
                    st.image(uploaded_image,
                             caption="Detected Image",
                             use_column_width=True)

#视频推理
def infer_uploaded_video(model):
    source_video = st.file_uploader(
        label="选择一个视频...",
        type=("mp4", "flv", "mov", 'avi',)
    )

    col1, col2 = st.columns(2)
    with col1:
        if source_video:
            st.video(source_video)

    if source_video:
        if st.button("Execution"):
            with st.spinner("执行中..."):
                try:
                    tfile = tempfile.NamedTemporaryFile()
                    tfile.write(source_video.read())
                    vid_cap = cv2.VideoCapture(
                        tfile.name)
                    st_frame = st.empty()
                    with col2:
                        while (vid_cap.isOpened()):
                            success, image = vid_cap.read()
                            if success:
                                _display_detected_frames(
                                    model,
                                    st_frame,
                                    image
                                )
                            else:
                                vid_cap.release()
                                break
                except Exception as e:
                     print(e)
                     st.error(f"Error loading video: {e}")

def infer_uploaded_webcam(model):

    try:
        flag = st.button(
            label="Stop running"
        )
        vid_cap = cv2.VideoCapture(0)  # 本地摄像头
        st_frame = st.empty()
        while not flag:
            success, image = vid_cap.read()
            if success:
                _display_detected_frames(
                    model,
                    st_frame,
                    image
                )
            else:
                vid_cap.release()
                break
    except Exception as e:
        st.error(f"Error loading video: {str(e)}")
